from typing import List, Dict, Any, Optional
from openai import OpenAI
from .base import BaseLLM
from ..server.config import get_settings

class MoonshotLLM(BaseLLM):
    def __init__(
        self,
        model: str = "moonshot-v1-8k",  # Moonshot的模型名称
        temperature: float = 0.7,
        max_tokens: int = 1000
    ):
        """
        初始化Moonshot LLM
        参数:
            model: 模型名称 (moonshot-v1-8k, moonshot-v1-32k, moonshot-v1-128k)
            temperature: 温度参数
            max_tokens: 最大生成token数
        """
        settings = get_settings()
        
        # 初始化客户端
        self.client = OpenAI(
            api_key=settings.MOONSHOT_API_KEY,
            base_url="https://api.moonshot.cn/v1"
        )
        
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
    
    def generate(
        self,
        prompt: str,
        context: Optional[List[Dict[str, Any]]] = None,
        **kwargs
    ) -> str:
        """
        生成回答
        """
        # 构建完整提示词
        if context:
            full_prompt = self._build_prompt(prompt, context)
        else:
            full_prompt = prompt
            
        try:
            # 调用Moonshot API
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是由Moonshot AI提供的人工智能助手。"},
                    {"role": "user", "content": full_prompt}
                ],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                **kwargs
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            error_msg = f"Moonshot API调用出错: {str(e)}"
            print(error_msg)
            return error_msg
